import uuid
from typing import Dict, Optional

from langchain.memory import ConversationBufferWindowMemory


class SessionMemoryManager:
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance.sessions = {}
        return cls._instance

    def get_memory(self, session_id: str) -> ConversationBufferWindowMemory:
        if session_id not in self.sessions:
            self.sessions[session_id] = ConversationBufferWindowMemory(k=6, return_messages=True)
        return self.sessions[session_id]

    def generate_session_id(self) -> str:
        session_id = str(uuid.uuid4())
        self.sessions[session_id] = ConversationBufferWindowMemory(k=6, return_messages=True)
        return session_id

    def clear_session(self, session_id: str):
        if session_id in self.sessions:
            del self.sessions[session_id]


# 全局实例
memory_manager = SessionMemoryManager()


def get_memory_manager() -> SessionMemoryManager:
    """获取记忆管理器实例
    
    Returns:
        记忆管理器实例
    """
    return memory_manager